#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用 kernel 中的 softirq tracepoint 跟踪线程被软中断打断的调度轨迹。

该脚本需要 root 权限来访问 tracefs。
"""

import argparse
import json
import os
import re
import signal
import subprocess
import sys
import time
from collections import defaultdict

# 全局变量用于信号处理
stop_tracing = False


def signal_handler(sig, frame):
    """处理 Ctrl+C 信号以优雅地停止跟踪。"""
    global stop_tracing
    print("\n收到中断信号，正在停止跟踪...")
    stop_tracing = True


def get_task_cpu_mask(pid):
    """
    使用 taskset -p <pid> 获取进程绑定的 CPU 列表，并转换为 tracing_cpumask 格式。

    :param pid: 目标进程 PID
    :return: (cpu_list, hex_mask_str) 元组，例如 ([0, 1], "0x3")
    """
    try:
        # 执行 taskset -p <pid> 命令
        result = subprocess.run(['taskset', '-p', str(pid)], capture_output=True, text=True, check=True)
        output = result.stdout.strip()
        # 示例输出: "pid 1234's current affinity mask: 3" 或 "pid 1234's current affinity list: 0,1"
        # 我们处理 "mask:" 格式，因为它更容易转换为十六进制
        if "affinity mask:" in output:
            # 提取十六进制掩码 (例如 "3")
            hex_mask_str = output.split("affinity mask:")[-1].strip()
            # 转换为整数以计算位数
            mask_int = int(hex_mask_str, 16)
            # 计算出被设置的 CPU 核心列表
            cpu_list = []
            bit_position = 0
            temp_mask = mask_int
            while temp_mask > 0:
                if temp_mask & 1:
                    cpu_list.append(bit_position)
                temp_mask >>= 1
                bit_position += 1
            return cpu_list, f"0x{hex_mask_str}"
        elif "affinity list:" in output:
            # 如果是列表格式 (例如 "0,1,3-5")，则需要转换
            list_str = output.split("affinity list:")[-1].strip()
            cpu_list = []
            for part in list_str.split(','):
                if '-' in part:
                    start, end = map(int, part.split('-'))
                    cpu_list.extend(range(start, end + 1))
                else:
                    cpu_list.append(int(part))
            
            # 将 CPU 列表转换为十六进制掩码
            max_cpu = max(cpu_list) if cpu_list else 0
            mask_int = 0
            for cpu in cpu_list:
                mask_int |= (1 << cpu)
            # 格式化为带 "0x" 前缀的小写十六进制字符串
            hex_mask_str = hex(mask_int)
            return sorted(cpu_list), hex_mask_str
        else:
             raise ValueError(f"无法从 taskset 输出解析 CPU 信息: {output}")

    except subprocess.CalledProcessError as e:
        print(f"错误：执行 'taskset -p {pid}' 失败: {e.stderr}")
        sys.exit(1)
    except ValueError as e:
        print(f"错误：解析 taskset 输出失败: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"错误：获取 CPU 绑定信息时发生未知错误: {e}")
        sys.exit(1)


def setup_tracing(pid, delay_us):
    """
    配置并启动内核跟踪。

    :param pid: 目标进程 PID
    :param delay_us: 延迟阈值（微秒）
    """
    tracefs_path = "/sys/kernel/debug/tracing"
    if not os.path.exists(tracefs_path):
        raise FileNotFoundError(f"Tracefs 未挂载在 {tracefs_path}。请运行 'mount -t debugfs none /sys/kernel/debug'")

    # --- 新增：获取 CPU 绑定信息 ---
    cpu_list, hex_mask = get_task_cpu_mask(pid)
    print(f"检测到 PID {pid} 绑定到 CPU: {cpu_list} (掩码: {hex_mask})")
    # --- 新增结束 ---

    # 清除之前的设置
    with open(f"{tracefs_path}/trace", 'w') as f:
        f.write("")

    # 初始化
    with open(f"{tracefs_path}/buffer_size_kb", 'w') as f:
        f.write("140800")
    with open(f"{tracefs_path}/current_tracer", 'w') as f:
        f.write("nop")
    with open(f"{tracefs_path}/tracing_on", 'w') as f:
        f.write("0")
    with open(f"{tracefs_path}/trace", 'w') as f:
        f.write("") # 再次清空，确保干净

    # --- 新增：设置 tracing_cpumask ---
    try:
        with open(f"{tracefs_path}/tracing_cpumask", 'w') as f:
            f.write(hex_mask)
        print(f"已设置 tracing_cpumask 为 {hex_mask}，仅跟踪 CPU {cpu_list}")
    except PermissionError:
        print("警告：写入 tracing_cpumask 失败。请确保以 root 权限运行此脚本。")
        # 不退出，因为过滤器可能仍然有效，但这会影响性能
    except IOError as e:
        print(f"警告：写入 tracing_cpumask 时出错: {e}")
        # 不退出
    # --- 新增结束 ---

    # 设置过滤器，只跟踪指定 PID 的事件
    filter_str = f'common_pid == {pid}'
    try:
        with open(f"{tracefs_path}/events/irq/softirq_entry/filter", 'w') as f:
            f.write(filter_str)
        with open(f"{tracefs_path}/events/irq/softirq_exit/filter", 'w') as f:
            f.write(filter_str)
    except PermissionError:
        print("错误：写入过滤器失败。请确保以 root 权限运行此脚本。")
        sys.exit(1)

    # 启用 softirq tracepoints
    with open(f"{tracefs_path}/events/irq/softirq_entry/enable", 'w') as f:
        f.write("1")
    with open(f"{tracefs_path}/events/irq/softirq_exit/enable", 'w') as f:
        f.write("1")	

    # 设置跟踪选项
    #options_to_set = ["overwrite", " tee"]
    #for option in options_to_set:
    #    try:
    #        with open(f"{tracefs_path}/options/{option}", 'w') as f:
    #            f.write("1")
    #    except IOError:
            # 某些选项可能不可用或已默认启用
    #        pass

    # 设置跟踪时钟
    with open(f"{tracefs_path}/trace_clock", 'w') as f:
        f.write("mono")

    # 开启trace
    with open(f"{tracefs_path}/tracing_on", 'w') as f:
        f.write("1")

    print(f"开始跟踪 PID {pid} 的软中断活动，延迟阈值: {delay_us / 1000.0:.2f} 毫秒...")


def stop_tracing_and_collect(tracefs_path):
    """
    停止跟踪并收集数据。

    :param tracefs_path: tracefs 路径
    :return: 包含跟踪数据的列表
    """
    # 禁用 tracepoints
    with open(f"{tracefs_path}/events/irq/softirq_entry/enable", 'w') as f:
        f.write("0")
    with open(f"{tracefs_path}/events/irq/softirq_exit/enable", 'w') as f:
        f.write("0")

    # 关闭trace
    #with open(f"{tracefs_path}/tracing_on", 'w') as f:
    #    f.write("0")

    # 读取跟踪数据
    with open(f"{tracefs_path}/trace", 'r') as f:
        trace_data = f.read()

    # 清除缓冲区
    with open(f"{tracefs_path}/trace", 'w') as f:
        f.write("")

    return trace_data


def parse_trace_data(trace_data, delay_threshold_us):
    """
    解析原始跟踪数据，提取超过阈值的软中断事件。

    :param trace_data: 原始跟踪数据字符串
    :param delay_threshold_us: 延迟阈值（微秒）
    :return: Chrome Trace Format 的事件列表
    """
    chrome_trace_events = []
    # 使用字典存储每个软中断向量的 entry 时间
    softirq_entries = defaultdict(dict)

    # 正则表达式匹配 softirq_entry 和 softirq_exit
    # 示例行:
    # <...>-12345 [001] .... 123456.789123: softirq_entry: vec=6 [action=TASKLET]
    # <...>-12345 [001] .... 123456.789456: softirq_exit: vec=6 [action=TASKLET]
    entry_pattern = re.compile(
        r'^.*-(\d+)\s+\[(\d\d\d)\].*\s(\d+)\.(\d+):\ssoftirq_entry:\s+vec=(\d+)\s\[action=(\w+)\]'
    )
    exit_pattern = re.compile(
        r'^.*-(\d+)\s+\[(\d\d\d)\].*\s(\d+)\.(\d+):\ssoftirq_exit:\s+vec=(\d+)\s\[action=(\w+)\]'
    )

    for line in trace_data.splitlines():
        entry_match = entry_pattern.match(line)
        exit_match = exit_pattern.match(line)

        if entry_match:
            pid, cpu, sec, usec, vec, action = entry_match.groups()
            timestamp_ns = int(sec) * 1_000_000_000 + int(usec) * 1_000
            key = (pid, cpu, vec, action)
            softirq_entries[key]['entry_ts'] = timestamp_ns

        elif exit_match:
            pid, cpu, sec, usec, vec, action = exit_match.groups()
            timestamp_ns = int(sec) * 1_000_000_000 + int(usec) * 1_000
            key = (pid, cpu, vec, action)

            if key in softirq_entries and 'entry_ts' in softirq_entries[key]:
                entry_ts = softirq_entries[key]['entry_ts']
                duration_ns = timestamp_ns - entry_ts
                duration_us = duration_ns / 1000.0

                if duration_us >= delay_threshold_us:
                    # 创建 Chrome Trace 事件
                    # ph: "X" 表示 Complete Event (有开始和结束时间)
                    # pid: 进程ID
                    # tid: 线程ID (这里用CPU核心表示)
                    # ts: 开始时间戳 (微秒)
                    # dur: 持续时间 (微秒)
                    # name: 事件名称
                    chrome_event = {
                        "name": f"softirq:{action}({vec})",
                        "ph": "X",
                        "ts": entry_ts / 1000.0, # Chrome Trace 使用微秒
                        "dur": duration_us,
                        "pid": int(pid),
                        "tid": int(cpu), # 使用CPU作为线程ID来区分不同核心上的事件
                        "args": {
                            "vector": int(vec),
                            "action": action,
                            "duration_us": round(duration_us, 2) # 保留两位小数
                        }
                    }
                    chrome_trace_events.append(chrome_event)
                
                # 从字典中移除已处理的 entry
                del softirq_entries[key]

    return chrome_trace_events


def main():
    """主函数"""
    global stop_tracing

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("PID", type=int, help="要跟踪的进程 PID")
    parser.add_argument("DELAY", type=float, help="记录轨迹的延迟阈值（毫秒）")
    parser.add_argument("-o", "--output", type=str, default="softirq_trace.json", help="输出 JSON 文件名 (默认: softirq_trace.json)")

    args = parser.parse_args()

    pid = args.PID
    delay_ms = args.DELAY
    delay_us = delay_ms * 1000  # 转换为微秒
    output_file = args.output

    tracefs_path = "/sys/kernel/debug/tracing"

    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        setup_tracing(pid, delay_us)
        
        # 等待事件或中断
        while not stop_tracing:
            time.sleep(1) # 每秒检查一次

    finally:
        print("正在收集并处理跟踪数据...")
        raw_trace_data = stop_tracing_and_collect(tracefs_path)
        events = parse_trace_data(raw_trace_data, delay_us)

        if events:
            # 按时间戳排序
            events.sort(key=lambda x: x['ts'])
            
            # Chrome Trace JSON 格式
            chrome_trace_format = {
                "traceEvents": events,
                "displayTimeUnit": "ms",
                "systemTraceEvents": "",
                "otherData": {
                    "version": "Performance Tracer 1.1 (with CPU mask)",
                    "pid": pid,
                    "delay_threshold_ms": delay_ms
                }
            }

            with open(output_file, 'w') as f:
                json.dump(chrome_trace_format, f, indent=4)

            print(f"跟踪完成。找到 {len(events)} 个超过 {delay_ms} ms 的软中断事件。")
            print(f"Chrome Trace 数据已保存至: {output_file}")
            print("您可以将此文件导入 Chrome 浏览器的 'chrome://tracing' 工具进行可视化分析。")
        else:
            print(f"跟踪完成。未找到超过 {delay_ms} ms 的软中断事件。")


if __name__ == "__main__":
    if os.geteuid() != 0:
        print("此脚本需要 root 权限来访问内核跟踪点和 taskset。请使用 'sudo' 运行。")
        sys.exit(1)
    main()



